#![allow(dead_code)]

use pest::error::LineColLocation;
use pest::iterators::{Pair, Pairs};
use pest::Parser;

use crate::lexer::{Rule, SysyParser};

fn write_indent(out: &mut String, indent: usize) {
    for _ in 0..indent {
        out.push(' ');
    }
}

fn fmt_comp_unit(pairs: Pairs<Rule>, out: &mut String) {
    let mut first = true;
    for pair in pairs {
        match pair.as_rule() {
            Rule::VarDecl => {
                if !first {
                    out.push('\n');
                }
                fmt_var_decl(pair, out, 0);
            }
            Rule::ConstDecl => {
                if !first {
                    out.push('\n');
                }
                fmt_const_decl(pair, out, 0);
            }
            Rule::FuncDef => {
                if !first {
                    out.push('\n');
                    out.push('\n');
                }
                fmt_func_def(pair, out, 0);
            }
            Rule::INVALID => {}
            _ => {}
        }
        first = false;
    }
}

fn fmt_btype(pair: Pair<Rule>, out: &mut String) {
    let token = pair.into_inner().next().unwrap();
    out.push_str(token.as_str());
}

fn fmt_const_init_val(pair: Pair<Rule>, out: &mut String) {
    let inner: Vec<_> = pair.into_inner().collect();
    if inner.is_empty() {
        return;
    }
    match inner[0].as_rule() {
        Rule::ConstExp => fmt_exp_like(inner[0].clone(), out),
        Rule::L_BRACE => {
            for p in inner {
                match p.as_rule() {
                    Rule::L_BRACE | Rule::R_BRACE => out.push_str(p.as_str()),
                    Rule::COMMA => out.push_str(", "),
                    Rule::ConstInitVal => fmt_const_init_val(p, out),
                    _ => {}
                }
            }
        }
        _ => {}
    }
}

fn fmt_const_def(pair: Pair<Rule>, out: &mut String) {
    let mut inner = pair.into_inner();
    let ident = inner.next().unwrap();
    out.push_str(ident.as_str());
    let mut after_assign = false;
    for p in inner {
        match p.as_rule() {
            Rule::L_BRACKT => out.push('['),
            Rule::R_BRACKT => out.push(']'),
            Rule::ConstExp | Rule::Exp => {
                if after_assign {
                    out.push_str(" = ");
                    after_assign = false;
                }
                fmt_exp_like(p, out);
            }
            Rule::ASSIGN => after_assign = true,
            Rule::L_BRACE => {
                if after_assign {
                    out.push_str(" = ");
                    after_assign = false;
                }
                out.push('{');
            }
            Rule::R_BRACE => out.push('}'),
            Rule::COMMA => out.push_str(", "),
            Rule::ConstInitVal => {
                if after_assign {
                    out.push_str(" = ");
                    after_assign = false;
                }
                fmt_const_init_val(p, out);
            }
            _ => {}
        }
    }
}

fn fmt_const_decl(pair: Pair<Rule>, out: &mut String, indent: usize) {
    let mut inner = pair.into_inner();
    write_indent(out, indent);
    let kw = inner.next().unwrap();
    debug_assert_eq!(kw.as_rule(), Rule::CONST);
    out.push_str("const ");
    let btype = inner.next().unwrap();
    fmt_btype(btype, out);
    out.push(' ');
    let mut first = true;
    for p in inner {
        match p.as_rule() {
            Rule::ConstDef => {
                if !first {
                    out.push_str(", ");
                }
                fmt_const_def(p, out);
                first = false;
            }
            Rule::SEMICOLON => out.push_str(";"),
            Rule::COMMA => {}
            _ => {}
        }
    }
}

fn fmt_var_def(pair: Pair<Rule>, out: &mut String) {
    let mut inner = pair.into_inner();
    let ident = inner.next().unwrap();
    out.push_str(ident.as_str());
    let mut after_assign = false;
    for p in inner {
        match p.as_rule() {
            Rule::L_BRACKT => out.push('['),
            Rule::R_BRACKT => out.push(']'),
            Rule::ConstExp | Rule::Exp => {
                if after_assign {
                    out.push_str(" = ");
                    after_assign = false;
                }
                fmt_exp_like(p, out);
            }
            Rule::ASSIGN => after_assign = true,
            Rule::L_BRACE => {
                if after_assign {
                    out.push_str(" = ");
                    after_assign = false;
                }
                out.push('{');
            }
            Rule::R_BRACE => out.push('}'),
            Rule::COMMA => out.push_str(", "),
            _ => {}
        }
    }
}

fn fmt_var_decl(pair: Pair<Rule>, out: &mut String, indent: usize) {
    let mut inner = pair.into_inner();
    write_indent(out, indent);
    let btype = inner.next().unwrap();
    fmt_btype(btype, out);
    out.push(' ');
    let mut first = true;
    for p in inner {
        match p.as_rule() {
            Rule::VarDef => {
                if !first {
                    out.push_str(", ");
                }
                fmt_var_def(p, out);
                first = false;
            }
            Rule::SEMICOLON => out.push_str(";"),
            Rule::COMMA => {}
            _ => {}
        }
    }
}

fn fmt_func_def(pair: Pair<Rule>, out: &mut String, indent: usize) {
    let mut inner = pair.into_inner();
    write_indent(out, indent);
    let func_type = inner.next().unwrap();
    out.push_str(func_type.as_str());
    out.push(' ');
    let ident = inner.next().unwrap();
    out.push_str(ident.as_str());
    let lparen = inner.next().unwrap();
    debug_assert_eq!(lparen.as_rule(), Rule::L_PAREN);
    out.push('(');
    let mut next = inner.next().unwrap();
    if next.as_rule() == Rule::FuncFParams {
        fmt_func_fparams(next, out);
        next = inner.next().unwrap();
    }
    debug_assert_eq!(next.as_rule(), Rule::R_PAREN);
    out.push(')');
    out.push(' ');
    let block = inner.next().unwrap();
    fmt_block(block, out, indent);
}

fn fmt_func_fparams(pair: Pair<Rule>, out: &mut String) {
    let mut first = true;
    for p in pair.into_inner() {
        match p.as_rule() {
            Rule::FuncFParam => {
                if !first {
                    out.push_str(", ");
                }
                fmt_func_fparam(p, out);
                first = false;
            }
            Rule::COMMA => {}
            _ => {}
        }
    }
}

fn fmt_func_fparam(pair: Pair<Rule>, out: &mut String) {
    let mut inner = pair.into_inner();
    let btype = inner.next().unwrap();
    fmt_btype(btype, out);
    out.push(' ');
    let ident = inner.next().unwrap();
    out.push_str(ident.as_str());
    for p in inner {
        match p.as_rule() {
            Rule::L_BRACKT | Rule::R_BRACKT => out.push_str(p.as_str()),
            Rule::ConstExp => fmt_exp_like(p, out),
            _ => {}
        }
    }
}

fn fmt_block(pair: Pair<Rule>, out: &mut String, indent: usize) {
    out.push_str("{\n");
    for p in pair.into_inner() {
        match p.as_rule() {
            Rule::Stmt => {
                fmt_stmt(p, out, indent + 4);
                out.push('\n');
            }
            Rule::VarDecl => {
                fmt_var_decl(p, out, indent + 4);
                out.push('\n');
            }
            Rule::ConstDecl => {
                fmt_const_decl(p, out, indent + 4);
                out.push('\n');
            }
            Rule::L_BRACE | Rule::R_BRACE => {}
            _ => {}
        }
    }
    write_indent(out, indent);
    out.push('}');
}

fn fmt_stmt(pair: Pair<Rule>, out: &mut String, indent: usize) {
    let inner = pair.into_inner().collect::<Vec<_>>();
    if inner.is_empty() {
        return;
    }
    match inner[0].as_rule() {
        Rule::LVal => {
            write_indent(out, indent);
            fmt_lval(inner[0].clone(), out);
            out.push_str(" = ");
            fmt_exp_like(inner[2].clone(), out);
            out.push_str(";");
        }
        Rule::Block => {
            write_indent(out, indent);
            fmt_block(inner[0].clone(), out, indent);
        }
        Rule::IF => {
            write_indent(out, indent);
            out.push_str("if (");
            fmt_cond(inner[2].clone(), out);
            out.push(')');
            let then_stmt = inner[4].clone();
            let then_first_rule = then_stmt.clone().into_inner().next().unwrap().as_rule();
            match then_first_rule {
                Rule::Block => {
                    out.push(' ');
                    fmt_block(then_stmt.into_inner().next().unwrap(), out, indent);
                }
                _ => {
                    out.push('\n');
                    fmt_stmt(then_stmt, out, indent + 4);
                }
            }
            if inner.len() > 5 {
                let mut opt_else = Some(inner[6].clone());
                while let Some(else_stmt) = opt_else {
                    let parts = else_stmt.clone().into_inner().collect::<Vec<_>>();
                    let first = parts[0].clone();
                    match first.as_rule() {
                        Rule::IF => {
                            out.push('\n');
                            write_indent(out, indent);
                            out.push_str("else if (");
                            fmt_cond(parts[2].clone(), out);
                            out.push(')');
                            let ei_then = parts[4].clone();
                            let ei_then_first_rule =
                                ei_then.clone().into_inner().next().unwrap().as_rule();
                            match ei_then_first_rule {
                                Rule::Block => {
                                    out.push(' ');
                                    fmt_block(ei_then.into_inner().next().unwrap(), out, indent);
                                }
                                _ => {
                                    out.push('\n');
                                    fmt_stmt(ei_then, out, indent + 4);
                                }
                            }
                            if parts.len() > 5 {
                                opt_else = Some(parts[6].clone());
                            } else {
                                opt_else = None;
                            }
                        }
                        Rule::Block => {
                            out.push('\n');
                            write_indent(out, indent);
                            out.push_str("else ");
                            fmt_block(first, out, indent);
                            opt_else = None;
                        }
                        _ => {
                            out.push('\n');
                            write_indent(out, indent);
                            out.push_str("else\n");
                            fmt_stmt(else_stmt, out, indent + 4);
                            opt_else = None;
                        }
                    }
                }
            }
        }
        Rule::WHILE => {
            write_indent(out, indent);
            out.push_str("while (");
            fmt_cond(inner[2].clone(), out);
            out.push(')');
            let body = inner[4].clone();
            let body_first_rule = body.clone().into_inner().next().unwrap().as_rule();
            match body_first_rule {
                Rule::Block => {
                    out.push(' ');
                    fmt_block(body.into_inner().next().unwrap(), out, indent);
                }
                _ => {
                    out.push('\n');
                    fmt_stmt(body, out, indent + 4);
                }
            }
        }
        Rule::BREAK | Rule::CONTINUE => {
            write_indent(out, indent);
            out.push_str(inner[0].as_str());
            out.push_str(";");
        }
        Rule::RETURN => {
            write_indent(out, indent);
            out.push_str("return");
            if inner.len() >= 3 && inner[1].as_rule() != Rule::SEMICOLON {
                out.push(' ');
                let exp_idx = if inner[1].as_rule() == Rule::Exp { 1 } else { 2 };
                fmt_exp_like(inner[exp_idx].clone(), out);
            }
            out.push_str(";");
        }
        Rule::SEMICOLON => {
            write_indent(out, indent);
            out.push_str(";");
        }
        Rule::Exp => {
            write_indent(out, indent);
            fmt_exp_like(inner[0].clone(), out);
            out.push_str(";");
        }
        _ => {}
    }
}

fn fmt_cond(pair: Pair<Rule>, out: &mut String) {
    fmt_lor_exp(pair.into_inner().next().unwrap(), out);
}

fn fmt_exp_like(pair: Pair<Rule>, out: &mut String) {
    match pair.as_rule() {
        Rule::ConstExp | Rule::Exp => fmt_add_exp(pair.into_inner().next().unwrap(), out),
        Rule::Number => out.push_str(pair.as_str()),
        _ => fmt_add_exp(pair, out),
    }
}

fn fmt_primary_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    match first.as_rule() {
        Rule::LVal => fmt_lval(first, out),
        Rule::Number => out.push_str(first.as_str()),
        Rule::L_PAREN => {
            out.push('(');
            let e = it.next().unwrap();
            fmt_exp_like(e, out);
            out.push(')');
        }
        _ => {}
    }
}

fn fmt_unary_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner().peekable();
    let first = it.peek().unwrap().as_rule();
    if matches!(first, Rule::PrimaryExp | Rule::IDENT) {
        let first_pair = it.next().unwrap();
        match first_pair.as_rule() {
            Rule::PrimaryExp => fmt_primary_exp(first_pair, out),
            Rule::IDENT => {
                out.push_str(first_pair.as_str());
                out.push('(');
                if let Some(_next) = it.next() {
                    if let Some(p) = it.peek() {
                        if p.as_rule() == Rule::FuncRParams {
                            fmt_func_rparams(it.next().unwrap().clone(), out);
                        }
                    }
                }
                out.push(')');
            }
            _ => {}
        }
    } else {
        let op = it.next().unwrap();
        let rhs = it.next().unwrap();
        let mut rhs_buf = String::new();
        fmt_unary_exp(rhs, &mut rhs_buf);
        while rhs_buf
            .chars()
            .next()
            .map(|c| c.is_whitespace())
            .unwrap_or(false)
        {
            rhs_buf.remove(0);
        }
        out.push_str(op.as_str());
        out.push_str(&rhs_buf);
    }
}

fn fmt_func_rparams(pair: Pair<Rule>, out: &mut String) {
    let mut first = true;
    for p in pair.into_inner() {
        match p.as_rule() {
            Rule::Exp => {
                if !first {
                    out.push_str(", ");
                }
                fmt_exp_like(p, out);
                first = false;
            }
            Rule::COMMA => {}
            _ => {}
        }
    }
}

fn fmt_mul_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_unary_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_unary_exp(rhs, out);
    }
}

fn fmt_add_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_mul_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_mul_exp(rhs, out);
    }
}

fn fmt_rel_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_add_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_add_exp(rhs, out);
    }
}

fn fmt_eq_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_rel_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_rel_exp(rhs, out);
    }
}

fn fmt_land_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_eq_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_eq_exp(rhs, out);
    }
}

fn fmt_lor_exp(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let first = it.next().unwrap();
    fmt_land_exp(first, out);
    while let Some(op) = it.next() {
        out.push(' ');
        out.push_str(op.as_str());
        out.push(' ');
        let rhs = it.next().unwrap();
        fmt_land_exp(rhs, out);
    }
}

fn fmt_lval(pair: Pair<Rule>, out: &mut String) {
    let mut it = pair.into_inner();
    let ident = it.next().unwrap();
    out.push_str(ident.as_str());
    for p in it {
        match p.as_rule() {
            Rule::L_BRACKT | Rule::R_BRACKT => out.push_str(p.as_str()),
            Rule::Exp => fmt_exp_like(p, out),
            _ => {}
        }
    }
}

fn collect_invalids(input: &str) -> Result<Vec<usize>, pest::error::Error<Rule>> {
    let mut pairs = SysyParser::parse(Rule::CompUnit, input)?;
    let comp = pairs.next().unwrap();
    fn walk(pair: Pair<Rule>, acc: &mut Vec<usize>) {
        match pair.as_rule() {
            Rule::INVALID
            | Rule::TOP_INVALID
            | Rule::PARAM_INVALID
            | Rule::ARGS_INVALID
            | Rule::INIT_INVALID
            | Rule::VARDEF_INVALID
            | Rule::CONSTDEF_INVALID => {
                acc.push(pair.line_col().0);
            }
            _ => {}
        }
        for c in pair.into_inner() {
            walk(c, acc);
        }
    }
    let mut acc = Vec::new();
    walk(comp, &mut acc);
    Ok(acc)
}

pub fn format(input: &str) -> String {
    match collect_invalids(input) {
        Ok(lines) => {
            if !lines.is_empty() {
                return lines
                    .into_iter()
                    .map(|line| format!(
                        "Error type B at Line {}: syntax error.",
                        line
                    ))
                    .collect::<Vec<_>>()
                    .join("\n");
            }
            let mut pairs = SysyParser::parse(Rule::CompUnit, input).expect("parse should succeed");
            let comp = pairs.next().unwrap();
            let mut out = String::new();
            fmt_comp_unit(comp.into_inner(), &mut out);
            out = out.replace("return - ", "return -");
            out = out.replace("return + ", "return +");
            out = out.replace("return ! ", "return !");
            out
        }
        Err(e) => {
            if let LineColLocation::Pos((line, _col)) = e.line_col {
                format!("Error type B at Line {}: {}", line, e.variant.message())
            } else {
                format!("Error type B at Line {}: {}", 1, e)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::fs;

    use super::format;

    #[test]
    fn test_file() {
        let filename = "tests/lab2-tests/test0.in";
        let input = fs::read_to_string(filename).expect("Failed to read file");
        let result = format(&input);
        assert!(!result.is_empty());
        println!("{}", result);
    }

    #[test]
    fn test_unary_and_binary_spaces() {
        let input = "int f(){int a=1;return -a + 2*-a;}";
        let out = format(input);
        let expect = "int f() {\n    int a = 1;\n    return -a + 2 * -a;\n}";
        assert_eq!(out, expect);
    }

    #[test]
    fn test_if_else_while_and_commas() {
        let input = "int f(int x,int y){if(x){x=x+1;}else if(y) x=x-1; else {while(x) x=x-1;} return x;}";
        let out = format(input);
        let expect = "int f(int x, int y) {\n    if (x) {\n        x = x + 1;\n    }\n    else if (y)\n        x = x - 1;\n    else {\n        while (x)\n            x = x - 1;\n    }\n    return x;\n}";
        assert_eq!(out, expect);
    }

    #[test]
    fn test_array_init_and_nested_braces() {
        let input = "int main(){int a[2]={1,2};int b[2][2]={{1,2},{3,4}};return 0;}";
        let out = format(input);
        let expect = "int main() {\n    int a[2] = {1, 2};\n    int b[2][2] = {{1, 2}, {3, 4}};\n    return 0;\n}";
        assert_eq!(out, expect);
    }

    #[test]
    fn test_const_decl_basic() {
        let input = "int f(){const int a=1;const int b[2]={1,2};return a+b[0];}";
        let out = format(input);
        let expect = "int f() {\n    const int a = 1;\n    const int b[2] = {1, 2};\n    return a + b[0];\n}";
        assert_eq!(out, expect);
    }
}
